import { createWorker } from "tesseract.js";
import * as tf from "@tensorflow/tfjs";
import { loadImg } from "../dataset";
import { loadGraphModel } from "@tensorflow/tfjs-converter";
import { CTCBeamSearch, Vocabulary } from "ctc-beam-search";
import content from "../char_std_all.txt";
let worker;
const target_height = 32;
const n_classes = 6510;
export async function initTesseract() {
  const start = Date.now();
  worker = createWorker({
    logger: (m) => console.log(m),
    // workerPath: ``
  });

  await worker.load();
  await worker.loadLanguage("chi_sim");
  await worker.initialize("chi_sim");
  console.warn("load tesseract duration: ", Date.now() - start, " ms");
  return 1;
}

export async function init() {
  const MODEL_URL = "http://localhost:9000/layerweb/model.json";
  const model = await tf.loadLayersModel(MODEL_URL);
  return model;
}

function beamSearch(data) {
  const words = content.split("\n").slice(1);
  const wordMap = words.reduce((acc, item, index) => {
    acc[item] = index;
    return acc;
  }, {});
  const bs = new CTCBeamSearch(new Vocabulary(wordMap, words.length - 1));
  const results = bs.search(data[0], 5);
  return results[0];
}

export async function recognize(modelPromise, imgData, draw) {
  const model = await modelPromise;
  imgData = await loadImg("/x.jpeg");
  // imgData = await imgData.data();
  // const [h, w] = imgData.shape;
  // const scale = h / target_height;

  const width = 280; //Math.floor(w / scale);
  imgData = tf.image.resizeBilinear(imgData, [target_height, width], true);
  imgData = tf.mean(imgData, 2).expandDims(-1);
  draw(imgData);
  imgData = tf.sub(tf.div(imgData, 255), 0.5);
  const predict = model.predict(imgData.expandDims(0));
  // console.log(a.dataSync());
  // const _predict = Array.prototype.slice.apply(
  //   tf.argMax(predict, 2).dataSync()
  // );

  const _predict = beamSearch(predict.mul(1e4).arraySync());
  // let chars = [];
  // _predict.forEach((item, i) => {
  //   if (
  //     item != n_classes - 1 &&
  //     !(
  //       (i > 0 && item === _predict[i - 1]) ||
  //       (i > 1 && item === _predict[i - 2])
  //     )
  //   ) {
  //     chars.push(item - 1);
  //   }
  // });
  // console.log(chars);
  if (1 == 1) return;
  model
    .executeAsync(
      {
        the_input: imgData.expandDims(0),
      },
      ["Identity:0"]
    )
    .then((predict) => {
      predict = Array.prototype.slice.apply(tf.argMax(predict, 2).dataSync());
      let chars = [];
      predict.forEach((item, i) => {
        if (
          item != n_classes - 1 &&
          !(
            (i > 0 && item === predict[i - 1]) ||
            (i > 1 && item === predict[i - 2])
          )
        ) {
          chars.push(item - 1);
        }
      });
      console.log(chars);
    });
  // await TesseractWorker;

  // const { data: { text } } = await worker.recognize(toImageFromBlob(imgData));
  // console.log(text);
  // await worker.terminate();
}

function toImageFromBlob(arrayBuffer) {
  var blob = new Blob(arrayBuffer, { type: "image/jpeg" });
  var urlCreator = window.URL || window.webkitURL;
  return urlCreator.createObjectURL(blob);
}
